import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeRegressor
from sklearn.metrics import mean_squared_error, r2_score

# 数据读取
data = pd.read_csv('data/data4.csv')

# 数据探索性分析
print('数据基本信息：')
data.info()

# 查看数据集行数和列数
rows, columns = data.shape

if rows < 1000:
    # 小数据集（行数少于 1000）查看全量数据信息
    print('数据全部内容信息：')
    print(data.to_csv(sep='\t', na_rep='nan'))
else:
    # 大数据集查看数据前几行信息
    print('数据前几行内容信息：')
    print(data.head().to_csv(sep='\t', na_rep='nan'))

# 数据清洗
# 处理缺失值
data['price'] = data['price'].fillna(data['price'].mean())
data['discount_rate'] = data['discount_rate'].fillna(data['discount_rate'].median())
data = data.dropna(subset=['sales_volume'])

# 特征工程
# 计算折扣后的价格
data['discounted_price'] = data['price'] * (1 - data['discount_rate'])

# 对商品类别进行独热编码
category_dummies = pd.get_dummies(data['category'], prefix='category')
data = pd.concat([data, category_dummies], axis=1)

# 特征选择
features = ['price', 'discount_rate', 'discounted_price'] + list(category_dummies.columns)
target = 'sales_volume'
X = data[features]
y = data[target]

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 模型训练
model = DecisionTreeRegressor(random_state=42)
model.fit(X_train, y_train)

# 模型预测
y_pred = model.predict(X_test)

# 模型评估
mse = mean_squared_error(y_test, y_pred)
rmse = np.sqrt(mse)
r2 = r2_score(y_test, y_pred)
print(f"均方误差 (MSE): {mse}")
print(f"均方根误差 (RMSE): {rmse}")
print(f"决定系数 (R²): {r2}")

# 数据可视化
# 不同商品类别的平均销量柱状图
category_avg_sales = data.groupby('category')['sales_volume'].mean()
plt.figure(figsize=(10, 6))
sns.barplot(x=category_avg_sales.index, y=category_avg_sales.values)
plt.title('不同商品类别的平均销量')
plt.xlabel('商品类别')
plt.xticks(rotation=45)
plt.ylabel('平均销量')
plt.show()

# 价格与销量的散点图
plt.figure(figsize=(10, 6))
plt.scatter(data['price'], data['sales_volume'])
plt.title('价格与销量的关系')
plt.xlabel('价格')
plt.ylabel('销量')
plt.show()